package org.ricks.agent;

import org.ricks.log.Logger;
import java.io.File;
import java.io.FileInputStream;
import java.lang.instrument.ClassDefinition;
import java.lang.instrument.Instrumentation;
import java.nio.file.Path;
import java.util.Arrays;

/**
 * @author chenwei
 * @Title:
 * @Package
 * @Description:
 * @date 2020/11/1717:24
 */
public class JavaAgent implements FileListener {

    private static Instrumentation inst;

    public static void premain(String args, Instrumentation inst){
        JavaAgent.inst = inst;
//        System.out.println("Hi, I'm premain!" + inst.toString());
    }

    public static void agentmain(String args, Instrumentation inst){
        JavaAgent.inst = inst;
//        System.out.println("Hi, I'm agentmain!" + inst.toString());
    }

    @Override
    public void onCreate(Path file) {
        loop(file.toFile());
    }

    @Override
    public void onModify(Path file) {
        loop(file.toFile());
    }

    @Override
    public void onDelete(Path path) {
        String absolutePath = path.toFile().getAbsolutePath();
        if (!absolutePath.endsWith(".class") && !absolutePath.endsWith(".dat")) path.toFile().mkdirs();
    }

    private void loop(File file){
        if(file.isDirectory()){
            Arrays.stream(file.listFiles()).forEach(f -> loop(f));
        } else redefineClasses(file);
    }

    private void redefineClasses(File classFile) {
        Logger.info("start load class file: " + classFile.getAbsolutePath());
        Arrays.stream(inst.getAllLoadedClasses()).filter(c -> {
            return classFile.getAbsolutePath().endsWith(c.getName().replaceAll("\\.","\\"+File.separator)+".class");
        }).forEach(c -> {
            try {
                byte[] bytes = fileToBytes(classFile);
                ClassDefinition classDefinition = new ClassDefinition(c,bytes);
                inst.redefineClasses(classDefinition);
            } catch (Exception e) {
                Logger.error("error redefined: {} ", classFile.getAbsolutePath());
                e.printStackTrace();
            }
            Logger.info("class load success: {} ", classFile.getAbsolutePath());
        });
    }

    private byte[] fileToBytes(File file) {
        byte[] bytes = new byte[(int) file.length()];
        try {
            FileInputStream in = new FileInputStream(file);
            in.read(bytes);
            in.close();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return bytes;
    }
}
